import asyncio
import uuid
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, Optional


from app.core.config import settings
from app.core.constants import TaskStatus, TriggerType
from app.core.task.executor.module_executor import ModuleExecutor
from app.core.task.executor.base_executor import BaseExecutor
from app.core.task.executor.script_executor import ScriptExecutor
from app.core.task.store.redis import RedisStore
from app.core.task.task_log import scheduler_logger
from app.schemas.job import SchedulerJob, RedisJobExecutionResult
from app.utils.distributed_lock import DistributedLock
from app.core.task.time_wheel import RedisTimeWheel
from app.utils.timezone import now
from app.core.redis_conn import redis_helper


@dataclass
class WorkerInfo:
    """Worker信息"""

    worker_id: str
    last_heartbeat: datetime
    status: str
    current_job_id: Optional[int] = None
    executor: Optional[BaseExecutor] = None


class Scheduler:
    _instance = None
    _initialized = False

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            raise RuntimeError("调度器尚未初始化")
        return cls._instance

    def __init__(self):
        if Scheduler._initialized:
            return
        self._connection = redis_helper.get_connection(settings.TASK_REDIS_DB)
        # 基础组件初始化
        self.running = False
        self.time_wheel = RedisTimeWheel()
        self.store = RedisStore()
        self.logger = scheduler_logger

        # Worker配置
        self.worker_num = settings.SCHEDULER_WORKER_NUM
        self.worker_timeout = settings.SCHEDULER_WORKER_TIMEOUT

        # 运行时状态
        self.workers: Dict[str, WorkerInfo] = {}
        self.task_queue = asyncio.PriorityQueue(maxsize=settings.SCHEDULER_MAX_QUEUE_SIZE)
        self.worker_lock_key = "scheduler:worker_lock"
        self._stop_event = asyncio.Event()

        Scheduler._initialized = True

    def _get_executor_for_job(self, job: SchedulerJob) -> BaseExecutor:
        """根据任务类型选择合适的执行器"""
        if job.executor_type == "script":
            return ScriptExecutor(self.store)
        return ModuleExecutor(self.store)

    async def _create_worker(self) -> str:
        """创建新的worker"""
        worker_id = f"worker-{uuid.uuid4()}"
        worker = WorkerInfo(
            worker_id=worker_id,
            last_heartbeat=datetime.now(),
            status="running",
            executor=None,  # 初始化时不创建执行器
        )
        self.workers[worker_id] = worker

        # 启动worker循环
        asyncio.create_task(self._worker_loop(worker))
        self.logger.info(f"创建worker: {worker_id}")
        return worker_id

    async def _worker_loop(self, worker: WorkerInfo):
        """Worker主循环"""
        try:
            while self.running and worker.status == "running":
                try:
                    # 获取任务
                    priority, job = await self.task_queue.get()
                    job = SchedulerJob.model_validate(job)
                    worker.current_job_id = job.id
                    worker.last_heartbeat = datetime.now()

                    try:
                        # 根据任务类型选择执行器
                        worker.executor = self._get_executor_for_job(job)
                        # 执行任务
                        execution_record = await worker.executor.execute_job(job)
                        # 处理执行结果
                        await self._handle_job_result(job, execution_record)

                    except asyncio.CancelledError:
                        self.logger.info(f"任务 {job.id} 被取消")
                        raise
                    except Exception as e:
                        self.logger.error(f"任务 {job.id} 执行失败: {str(e)}", exc_info=True)
                    finally:
                        worker.current_job_id = None
                        worker.executor = None  # 清理执行器
                        self.task_queue.task_done()
                        worker.last_heartbeat = datetime.now()

                except asyncio.CancelledError:
                    break
                except Exception as e:
                    self.logger.error(f"Worker {worker.worker_id} 执行出错: {str(e)}", exc_info=True)
                    await asyncio.sleep(1)
        finally:
            # 清理worker资源
            if worker.worker_id in self.workers:
                del self.workers[worker.worker_id]
            self.logger.info(f"Worker {worker.worker_id} 已停止")

    async def _handle_job_result(self, job: SchedulerJob, execution_record: RedisJobExecutionResult):
        """处理任务执行结果"""
        try:
            # 保存执行记录
            await self.store.add_job_execution_log(job.id, execution_record)

            # 处理循环任务
            if job.status == TaskStatus.ACTIVE and job.trigger_type in [TriggerType.INTERVAL, TriggerType.CRON]:
                next_run_time = await job.trigger.get_next_fire_time(now())
                if next_run_time:
                    job.next_run_time = next_run_time
                    await self.store.update_job(job.id, job)
                    await self.time_wheel.add_task(job.id, next_run_time)
        except Exception as e:
            self.logger.error(f"处理任务 {job.id} 结果失败: {str(e)}", exc_info=True)

    async def _monitor_workers(self):
        """监控worker状态"""
        while self.running:
            try:
                now = datetime.now()
                dead_workers = []

                # 检查worker心跳
                for worker_id, worker in self.workers.items():
                    if (now - worker.last_heartbeat).seconds > self.worker_timeout:
                        dead_workers.append(worker_id)

                # 移除死亡的worker
                for worker_id in dead_workers:
                    worker = self.workers.pop(worker_id)
                    self.logger.warning(f"Worker {worker_id} 已超时, " f"当前任务: {worker.current_job_id}")

                # 确保worker数量
                async with DistributedLock(self._connection, self.worker_lock_key):
                    while len(self.workers) < self.worker_num:
                        await self._create_worker()

            except Exception as e:
                self.logger.error(f"监控worker出错: {str(e)}", exc_info=True)

            await asyncio.sleep(self.worker_timeout / 2)

    async def start(self):
        """启动调度器"""
        if self.running:
            return

        self.running = True
        self._stop_event.clear()

        # 启动worker监控
        monitor_task = asyncio.create_task(self._monitor_workers())

        # 主调度循环
        try:
            while self.running:
                try:
                    # 获取到期的任务
                    due_tasks = await self.time_wheel.get_due_tasks()

                    # 将到期任务加入队列
                    # priority 越大，优先级越高

                    jobs = []
                    task_ids = [task_id for task_id, _ in due_tasks]
                    jobs = await self.store.get_jobs(task_ids)
                    jobs = sorted(jobs, key=lambda x: x.priority, reverse=True)

                    for job in jobs:
                        if job and job.status == TaskStatus.ACTIVE.value:
                            try:
                                # 使用优先级作为排序依据
                                await self.task_queue.put((-job.priority, job))
                            except asyncio.QueueFull:
                                self.logger.error(f"任务队列已满，无法添加任务: {job.id}")

                    await self.time_wheel.advance()

                except Exception as e:
                    self.logger.error(f"调度循环出错: {str(e)}", exc_info=True)

                await asyncio.sleep(0.1)  # 避免CPU过度使用

        finally:
            # 停止所有worker
            self.running = False
            monitor_task.cancel()
            try:
                await monitor_task
            except asyncio.CancelledError:
                pass

    async def shutdown(self):
        """关闭调度器"""
        if not self.running:
            return

        self.running = False
        self._stop_event.set()

        # 等待所有任务完成
        if self.task_queue.qsize() > 0:
            self.logger.info(f"等待 {self.task_queue.qsize()} 个任务完成...")
            await self.task_queue.join()

        # 清理workers
        for worker in list(self.workers.values()):
            worker.status = "stopped"
            if worker.current_job_id:
                try:
                    await worker.executor.cancel_job(worker.current_job_id)
                except Exception as e:
                    self.logger.error(f"取消任务 {worker.current_job_id} 失败: {str(e)}", exc_info=True)

        self.workers.clear()
        self.logger.info("调度器已关闭")

    async def cancel_job(self, job_id: str) -> bool:
        """取消指定任务"""
        for worker in self.workers.values():
            if worker.current_job_id == job_id:
                try:
                    return await worker.executor.cancel_job(job_id)
                except Exception as e:
                    self.logger.error(f"取消任务 {job_id} 失败: {str(e)}", exc_info=True)
                    return False
        return False

    async def scale_workers(self, new_count: int):
        """动态调整worker数量"""
        if new_count < 1:
            raise ValueError("worker数量必须大于0")

        async with DistributedLock(self._connection, self.worker_lock_key):
            self.worker_num = new_count
            current = len(self.workers)

            if new_count > current:
                # 增加workers
                for _ in range(new_count - current):
                    await self._create_worker()
            else:
                # 减少workers
                workers_to_remove = list(self.workers.keys())[new_count:]
                for worker_id in workers_to_remove:
                    worker = self.workers.pop(worker_id)
                    worker.status = "stopped"
                    if worker.current_job_id:
                        await worker.executor.cancel_job(worker.current_job_id)
                    self.logger.info(f"停止worker: {worker_id}")

    async def pause_job(self, job_id: str) -> bool:
        """暂停指定任务

        Returns:
            bool: 是否成功暂停
        """
        try:
            # 1. 更新存储中的任务状态
            job = await self.store.get_job(job_id)
            if not job:
                return False

            job.status = TaskStatus.PAUSED
            await self.store.update_job(job_id, job)

            # 2. 从时间轮中移除
            await self.time_wheel.remove_task(job_id)

            # 3. 如果任务正在执行，尝试取消
            await self.cancel_job(job_id)

            self.logger.info(f"任务已暂停: {job_id}")
            return True
        except Exception as e:
            self.logger.error(f"暂停任务 {job_id} 失败: {str(e)}", exc_info=True)
            return False

    async def resume_job(self, job_id: str) -> bool:
        """恢复指定任务

        Returns:
            bool: 是否成功恢复
        """
        try:
            # 1. 获取并更新任务状态
            job = await self.store.get_job(job_id)
            if not job:
                return False

            job.status = TaskStatus.ACTIVE

            # 2. 计算下次执行时间
            next_run_time = await job.trigger.get_next_fire_time(now())
            if next_run_time:
                job.next_run_time = next_run_time
                await self.store.update_job(job_id, job)
                await self.time_wheel.add_task(job_id, next_run_time)

            self.logger.info(f"任务已恢复: {job_id}")
            return True
        except Exception as e:
            self.logger.error(f"恢复任务 {job_id} 失败: {str(e)}", exc_info=True)
            return False

    async def delete_job(self, job_id: str) -> bool:
        """删除指定任务

        Returns:
            bool: 是否成功删除
        """
        try:
            # 1. 从存储中删除
            await self.store.delete_job(job_id)

            # 2. 从时间轮中移除
            await self.time_wheel.remove_task(job_id)

            # 3. 如果任务正在执行，尝试取消
            await self.cancel_job(job_id)

            self.logger.info(f"任务已删除: {job_id}")
            return True
        except Exception as e:
            self.logger.error(f"删除任务 {job_id} 失败: {str(e)}", exc_info=True)
            return False

    async def exists_job(self, job_id: str) -> bool:
        return await self.store.exists_job(job_id)

    async def update_job(self, job_id: str, job: SchedulerJob) -> bool:
        """更新指定任务

        Args:
            job_id: 任务ID
            job: 更新后的任务信息

        Returns:
            bool: 是否成功更新
        """
        try:
            # 1. 停止并取消当前任务
            await self.cancel_job(job_id)
            await self.time_wheel.remove_task(job_id)

            # 2. 更新任务信息
            if not await self.store.update_job(job_id, job):
                return False

            # 3. 如果任务处于激活状态，重新加入调度
            if job.status == TaskStatus.ACTIVE:
                # 计算下次执行时间
                next_run_time = await job.trigger.get_next_fire_time(now())
                if next_run_time:
                    job.next_run_time = next_run_time
                    await self.store.update_job(job_id, job)
                    await self.time_wheel.add_task(job_id, next_run_time)

            self.logger.info(f"任务已更新: {job_id}")
            return True

        except Exception as e:
            self.logger.error(f"更新任务 {job_id} 失败: {str(e)}", exc_info=True)
            return False

    async def add_job(self, job: SchedulerJob) -> bool:
        """添加任务

        Returns:
            bool: 是否成功添加
        """
        try:
            # 1. 添加任务
            await self.store.add_job(job)

            # 2. 启动任务
            if job.status == TaskStatus.ACTIVE:
                next_run_time = await job.trigger.get_next_fire_time(now())
                if next_run_time:
                    await self.time_wheel.add_task(job.id, next_run_time)

            self.logger.info(f"任务已添加: {job.id}")
            return True
        except Exception as e:
            self.logger.error(f"添加任务失败: {str(e)}", exc_info=True)
            return False
